import numpy as np


class CropTarget:
    """
    根据目标检测的box结果裁切成小图
    args:
        in:
            keys: 指定了从data中读取数据的关键字和保存数据的关键字
                格式: keys={
                    "det":      "目标检测结果的key",
                    "img_data": "原始输入数据的key"
                    "out":      "your_save_keyword"
                            }
            target_cls: 指定需要crop的目标分类
            expand_ratio: 裁切时需要向外延展一定比例, 确保完整的包含目标
        return:
            data[out]:  输出的结果会以keys["out]为关键字保存在data中
                        格式:
                        {
                        "imgs":[],         保存了扣下的图片
                        "clses":[],        对应扣下图片的类别id
                        "det_ids":[],      对应口下图片在目标检测结果det中的下标
                        "boxes":[],        经过expand后图片的boxes信息
                        }
    """

    def __init__(self, keys, target_cls, expand_ratio=0.2) -> None:
        self.keys = keys
        assert self.keys.get("img_data", False)
        self.target_cls = target_cls
        self.expand_ratio = expand_ratio
        pass

    def __call__(self, data):
        """将需要进行子任务的object裁出
        return:
            data: {img_data, ...}
        """
        det_res = data[self.keys["det"]]
        batch_input = data[self.keys["img_data"]]
        batch_boxes = np.array(det_res["boxes"])
        crop_res = []
        for input, boxes in zip(batch_input, batch_boxes):
            res_data = {
                "imgs": [],
                "clses": [],
                "det_ids": [],
                "boxes": [],
            }
            if len(boxes) == 0:
                crop_res.append(res_data)
                continue
            for idx_, box in enumerate(boxes):
                if box[-6] in self.target_cls:
                    crop_image, new_box = self.expand_crop(
                        input, box, self.expand_ratio
                    )  # 根据box大小扩展一圈后裁切
                    if crop_image is not None:
                        res_data["imgs"].append(crop_image)
                        res_data["det_ids"].append(idx_)
                        res_data["clses"].append(box[-6])
                        res_data["boxes"].append(new_box)

            crop_res.append(res_data)
        data[self.keys["out"]] = crop_res
        return data

    @staticmethod
    def expand_crop(images, rect, expand_ratio=0.2):
        imgh, imgw, c = images.shape
        _, _, xmin, ymin, xmax, ymax = [int(x) for x in rect[-6:]]  # rect[-6:].tolist()
        if expand_ratio != 0:
            h_half = (ymax - ymin) * (1 + expand_ratio) / 2.0
            w_half = (xmax - xmin) * (1 + expand_ratio) / 2.0
            center = [(ymin + ymax) / 2.0, (xmin + xmax) / 2.0]
            ymin = max(0, int(center[0] - h_half))
            ymax = min(imgh - 1, int(center[0] + h_half))
            xmin = max(0, int(center[1] - w_half))
            xmax = min(imgw - 1, int(center[1] + w_half))
        return images[ymin:ymax, xmin:xmax, :], [xmin, ymin, xmax, ymax]
